import os

import tensorflow as tf

tfrecord_dir = '/Users/d/Project/emotions/data/iemocap_4emo_spectr_norm_simple'

train_tfrecord = 'train.tfrecords'
vali_tfrecord = 'vali.tfrecords'
test_tfrecord = 'test.tfrecords'

count = 0


def parse_single_example(example_proto):
    features = {'data': tf.FixedLenFeature([], tf.string),
                'len': tf.FixedLenFeature([], tf.int64),
                'sentence_id': tf.FixedLenFeature([], tf.string),
                'label': tf.FixedLenFeature([], tf.int64)}
    parsed = tf.parse_single_example(example_proto, features=features)
    sentence_str = parsed['sentence_id']
    length = tf.cast(parsed['len'], tf.int32)
    label = tf.cast(parsed['label'], tf.int32)
    data = tf.reshape(tf.decode_raw(parsed['data'], tf.float32), [length, -1])
    return {'data': data, 'len': length, 'sentence_id': sentence_str,
            'label': label}


def get_dataset(tfrecord_path):
    dataset = tf.data.TFRecordDataset(tfrecord_path)
    dataset = dataset.map(parse_single_example)

    # dataset = dataset.shuffle(2)
    # dataset = dataset.batch(1)
    # dataset = dataset.padded_batch(2, padded_shapes={'data': [None, None],
    #                                                  'len': (),
    #                                                  'label': ()})
    dataset = dataset.padded_batch(20, padded_shapes={'data': [None, None],
                                                      'len': {},
                                                      'sentence_id': {},
                                                      'label': {}})
    return dataset


def main():
    global count
    tf_record_f = os.path.join(tfrecord_dir, test_tfrecord)
    data_set = get_dataset(tf_record_f)
    # data_set = tf.data.TFRecordDataset(tf_record_f)
    iterator = data_set.make_one_shot_iterator()
    next_item = iterator.get_next()
    loop_num = 10
    with tf.Session() as sess:
        # for i in range(loop_num):
        try:
            t = sess.run(next_item)
        except tf.errors.OutOfRangeError:
            t = None
        while t:
            # print(t['data'].shape)
            # print(np.mean(t['data'], axis=2))
            # print(np.std(t['data'], axis=2))
            # print(t['len'])
            # print(t['sentence_id'])
            print(t['label'])
            count += 1
            print(count)
            print('\n')
            try:
                t = sess.run(next_item)
            except tf.errors.OutOfRangeError:
                t = None


if __name__ == '__main__':
    main()
